from langchain_core.output_parsers import JsonOutputParser
from pydantic import BaseModel,Field
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import json
from model_utils import getLLM

class SubTitle(BaseModel):
    sub_titles: list[str] = Field("二级标题")

class Title(BaseModel):
    name : str = Field("一级标题")
    sub_titles:list[SubTitle] = Field("二级标题")

class Scheme(BaseModel):
    abstract:str = Field(description="摘要")
    outline:list[Title] = Field(description="大纲")

_system_prompt = """
您是一名优秀的方案规划师，您要根据提出的项目标题，规划一份完整的方案 大纲目录。
将大纲规划成两级目录，不要省略，二级目录3项。目录不需要序号。
还需要写份摘要，以便为后续的工作作为参考(无需其它多余信息)。
{output_template}
"""
_human_prompt = """
项目标题：{project_name}
"""

def create_planner(llm):
    prompt_template = ChatPromptTemplate.from_messages([
        ("system",_system_prompt),
        ("human",_human_prompt)
    ])

    parser = JsonOutputParser(pydantic_object=Scheme)
    prompt_template = prompt_template.partial(output_template=parser.get_format_instructions())

    return prompt_template | llm | parser

if __name__ == '__main__':
    llm = getLLM()
    agent = create_planner(llm)
    _reponese = agent.invoke({"project_name": "基于大模型的数字人驱动方案"})

    print(_reponese)